package org.xx.armory.commons;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.crypto.Cipher;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.UncheckedIOException;
import java.nio.charset.Charset;
import java.security.GeneralSecurityException;
import java.security.InvalidKeyException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Base64;

import static java.nio.charset.StandardCharsets.UTF_8;
import static org.xx.armory.commons.Validators.rejectIfNull;

/**
 * 用于加密和解密的工具类。
 */
public final class SecurityUtils {
    private static final String ENCRYPT_DECRYPT_ALGORITHM = "DES/ECB/PKCS5Padding";

    @SuppressWarnings("unused")
    private static final Logger LOGGER = LoggerFactory.getLogger(SecurityUtils.class);

    /**
     * 加密算法使用的随机数。
     * <p>{@link SecureRandom}类型是线程安全的，所以可以共用。</p>
     */
    private static final SecureRandom RANDOM = new SecureRandom();

    private SecurityUtils() {
        throw new AssertionError();
    }

    /**
     * 对字符串按照{@literal UTF_8}编码，然后加密, 并将加密结果编码为Base64格式
     *
     * @param s
     *         待加密的字符串
     * @param key
     *         用于加密的密钥。
     * @return 已加密的字符串
     * @throws UncheckedSecurityException
     *         如果加密过程中出现错误
     * @throws UncheckedIOException
     *         如果将字符串转为字节流或者将字节流转为字符串时出现错误。
     * @see #encryptWithBase64(String, Charset, SecretKey)
     */
    public static String encryptWithBase64(
            String s,
            SecretKey key
    ) {
        return encryptWithBase64(s, UTF_8, key);
    }

    /**
     * 对Base64格式字符串进行解密, 并将解密的结果按照{@literal UTF_8}编码为字符串
     *
     * @param s
     *         待解密的Base64格式字符串
     * @param key
     *         用于解密的密钥。
     * @return 已解密的字符串
     * @throws UncheckedSecurityException
     *         如果加密过程中出现错误
     * @throws UncheckedIOException
     *         如果将字符串转为字节流或者将字节流转为字符串时出现错误。
     * @see #decryptWithBase64(String, Charset, SecretKey)
     */
    public static String decryptWithBase64(
            String s,
            SecretKey key
    ) {
        return decryptWithBase64(s, UTF_8, key);
    }

    /**
     * 对字符串按照指定的编码格式编码，然后加密, 并将加密结果编码为Base64格式
     *
     * @param s
     *         待加密的字符串
     * @param charset
     *         将文本转化为字节的编码。
     * @param key
     *         用于加密的密钥。
     * @return 已加密的字符串
     * @throws UncheckedSecurityException
     *         如果加密过程中出现错误
     * @throws UncheckedIOException
     *         如果将字符串转为字节流或者将字节流转为字符串时出现错误。
     * @see #decryptWithBase64(String, Charset, SecretKey)
     * @see #encrypt(InputStream, OutputStream, SecretKey)
     */
    public static String encryptWithBase64(
            String s,
            Charset charset,
            SecretKey key
    ) {
        if (s == null || s.isEmpty()) {
            return "";
        } else {
            try (final ByteArrayInputStream src = new ByteArrayInputStream(s.getBytes(charset));
                 final ByteArrayOutputStream dst = new ByteArrayOutputStream()) {
                encrypt(src, dst, key);
                return Base64.getEncoder().encodeToString(dst.toByteArray());
            } catch (IOException ex) {
                throw new UncheckedIOException(ex);
            }
        }
    }

    /**
     * 对Base64格式字符串进行解密, 并将解密的结果按照指定编码为字符串
     *
     * @param s
     *         待解密的Base64格式字符串
     * @param charset
     *         将文本转化为字节的编码。
     * @param key
     *         用于解密的密钥。
     * @return 已解密的字符串
     * @throws UncheckedSecurityException
     *         如果解密过程中出现错误。
     * @throws UncheckedIOException
     *         如果将字符串转为字节流或者将字节流转为字符串时出现错误。
     * @see #encryptWithBase64(String, Charset, SecretKey)
     * @see #decrypt(InputStream, OutputStream, SecretKey)
     */
    public static String decryptWithBase64(
            String s,
            Charset charset,
            SecretKey key
    ) {
        if (s == null || s.isEmpty()) {
            return "";
        } else {
            final byte[] bytes = Base64.getDecoder().decode(s);
            try (final ByteArrayInputStream src = new ByteArrayInputStream(bytes);
                 final ByteArrayOutputStream dst = new ByteArrayOutputStream()) {
                decrypt(src, dst, key);
                return dst.toString(charset);
            } catch (IOException ex) {
                throw new UncheckedIOException(ex);
            }
        }
    }

    /**
     * 使用密钥对指定的输入流进行加密, 并将结果写入指定的输出流。
     *
     * <p>对于DES算法，默认模式为{@literal ECB/PCKS5Padding}</p>
     *
     * @param src
     *         指定的输入流。
     * @param dst
     *         指定的输出流。
     * @param key
     *         用于加密的密钥。
     * @throws UncheckedIOException
     *         如果读取流或者写入流过程中出现错误。
     * @throws IllegalStateException
     *         如果加密算法不存在。
     * @throws UncheckedSecurityException
     *         如果加密过程中出现错误
     * @throws IllegalArgumentException
     *         如果参数{@code src}是{@code null}或者{@code dst}或者{@code key}是{@code null}或者如果参数{@code key}表示不正确的密钥。
     */
    public static void encrypt(
            InputStream src,
            OutputStream dst,
            SecretKey key
    ) {
        rejectIfNull(src, "src");
        rejectIfNull(dst, "dst");
        rejectIfNull(key, "key");

        final int BUFFER_SIZE = 128;

        final Cipher cipher;
        try {
            cipher = Cipher.getInstance(ENCRYPT_DECRYPT_ALGORITHM);
            cipher.init(Cipher.ENCRYPT_MODE, key, RANDOM);
        } catch (NoSuchAlgorithmException | NoSuchPaddingException ex) {
            throw new IllegalStateException(ex);
        } catch (InvalidKeyException ex) {
            throw new IllegalArgumentException("illegal key", ex);
        }

        try {
            final byte[] buf = new byte[BUFFER_SIZE];
            int rl = src.read(buf);
            while (rl != -1) {
                dst.write(cipher.update(buf, 0, rl));
                rl = src.read(buf);
            }

            try {
                dst.write(cipher.doFinal());
            } catch (IOException ex) {
                throw new UncheckedIOException(ex);
            } catch (GeneralSecurityException ex) {
                throw new UncheckedSecurityException(ex);
            }

            dst.flush();
        } catch (IOException ex) {
            throw new UncheckedIOException(ex);
        }
    }

    /**
     * 使用唯一密钥对指定的输入流进行解密, 并将结果写入指定的输出流。
     *
     * @param src
     *         指定的输入流。
     * @param dst
     *         指定的输出流。
     * @param key
     *         用于解密的密钥。
     * @throws UncheckedIOException
     *         如果读取流或者写入流过程中出现错误。
     * @throws IllegalStateException
     *         如果加密算法不存在。
     * @throws UncheckedSecurityException
     *         如果解密过程中出现错误。
     * @throws IllegalArgumentException
     *         如果参数{@code src}是{@code null}或者{@code dst}或者{@code key}是{@code null}或者如果参数{@code key}表示不正确的密钥。
     */
    public static void decrypt(
            InputStream src,
            OutputStream dst,
            SecretKey key
    ) {
        rejectIfNull(src, "src");
        rejectIfNull(dst, "dst");
        rejectIfNull(key, "key");

        final int BUFFER_SIZE = 128;

        final Cipher cipher;
        try {
            cipher = Cipher.getInstance(ENCRYPT_DECRYPT_ALGORITHM);
            cipher.init(Cipher.DECRYPT_MODE, key, RANDOM);
        } catch (NoSuchAlgorithmException | NoSuchPaddingException ex) {
            throw new IllegalStateException(ex);
        } catch (InvalidKeyException ex) {
            throw new IllegalArgumentException("illegal key", ex);
        }

        try {
            final byte[] buf = new byte[BUFFER_SIZE];
            int rl = src.read(buf);
            while (rl != -1) {
                dst.write(cipher.update(buf, 0, rl));
                rl = src.read(buf);
            }

            try {
                dst.write(cipher.doFinal());
            } catch (GeneralSecurityException ex) {
                throw new UncheckedSecurityException(ex);
            }

            dst.flush();
        } catch (IOException ex) {
            throw new UncheckedIOException(ex);
        }
    }

    public static byte[] hash(
            String algorithm,
            byte[] src
    ) {
        final MessageDigest md;
        try {
            md = MessageDigest.getInstance(algorithm);
        } catch (GeneralSecurityException ex) {
            throw new IllegalStateException(ex);
        }

        if (src != null) {
            md.update(src);
            return md.digest();
        } else {
            return md.digest(new byte[0]);
        }
    }

    /**
     * 使用{@literal MD5}算法对指定的字节数组进行哈希。
     *
     * @param src
     *         待哈希的字节数组
     * @return 哈希的结果, 如果参数{@code src}是{@code null}, 那么相当于对长度为0的字节数组进行哈希。
     */
    public static byte[] md5(
            byte[] src
    ) {
        return hash("MD5", src);
    }

    /**
     * 使用{@literal SHA-1}算法对指定的字节数组进行哈希。
     *
     * @param src
     *         待哈希的字节数组
     * @return 哈希的结果, 如果参数{@code src}是{@code null}, 那么相当于对长度为0的字节数组进行哈希。
     */
    public static byte[] sha1(
            byte[] src
    ) {
        return hash("SHA-1", src);
    }

    /**
     * 使用{@literal SHA-256}算法对指定的字节数组进行哈希。
     *
     * @param src
     *         待哈希的字节数组
     * @return 哈希的结果, 如果参数{@code src}是{@code null}, 那么相当于对长度为0的字节数组进行哈希。
     */
    public static byte[] sha256(
            byte[] src
    ) {
        return hash("SHA-256", src);
    }

    /**
     * 使用{@literal SHA-1}算法对指定的长整数进行哈希。
     *
     * @param l
     *         待哈希的长整数。
     * @return 哈希的结果。
     * @see #sha1(byte[])
     */
    public static byte[] sha1(
            long l
    ) {
        final byte[] src = new byte[]{
                (byte) (l & 0xFF), (byte) ((l >> 8) & 0xFF),
                (byte) ((l >> 16) & 0xFF), (byte) ((l >> 24) & 0xFF),
                (byte) ((l >> 32) & 0xFF), (byte) ((l >> 40) & 0xFF),
                (byte) ((l >> 48) & 0xFF), (byte) ((l >> 56) & 0xFF),};

        return sha1(src);
    }

    /**
     * 使用{@literal SHA-1}算法对指定的整数进行哈希。
     *
     * @param i
     *         待哈希的整数。
     * @return 哈希的结果。
     * @see #sha1(byte[])
     */
    public static byte[] sha1(
            int i
    ) {
        final byte[] src = new byte[]{
                (byte) (i & 0xFF), (byte) ((i >> 8) & 0xFF),
                (byte) ((i >> 16) & 0xFF), (byte) ((i >> 24) & 0xFF)};

        return sha1(src);
    }

    /**
     * 使用{@literal MD5}算法对指定的字符串进行哈希。
     * <p>
     * 首先将字符串按照参数{@code charset}转化为字节数组, 然后进行哈希。
     * </p>
     *
     * @param s
     *         待哈希的字符串。
     * @param charset
     *         字符串编码。
     * @return 哈希的结果。
     * @throws IllegalArgumentException
     *         如果参数{@code charset}是{@code null}。
     * @see String#getBytes()
     * @see #md5(byte[])
     */
    public static byte[] md5(
            String s,
            Charset charset
    ) {
        rejectIfNull(charset, "charset");

        if (s == null) {
            return md5((byte[]) null);
        }

        return md5(s.getBytes(charset));
    }

    /**
     * 使用{@literal SHA-1}算法对指定的字符串进行哈希。
     * <p>
     * 首先将字符串按照参数{@code charset}转化为字节数组, 然后进行哈希。
     * </p>
     *
     * @param s
     *         待哈希的字符串。
     * @param charset
     *         字符串编码。
     * @return 哈希的结果。
     * @throws IllegalArgumentException
     *         如果参数{@code charset}是{@code null}。
     * @see String#getBytes()
     * @see #sha1(byte[])
     */
    public static byte[] sha1(
            String s,
            Charset charset
    ) {
        rejectIfNull(charset, "charset");

        if (s == null) {
            return sha1((byte[]) null);
        }

        return sha1(s.getBytes(charset));
    }

    /**
     * 使用{@literal SHA-256}算法对指定的字符串进行哈希。
     * <p>
     * 首先将字符串按照参数{@code charset}转化为字节数组, 然后进行哈希。
     * </p>
     *
     * @param s
     *         待哈希的字符串。
     * @param charset
     *         字符串编码。
     * @return 哈希的结果。
     * @throws IllegalArgumentException
     *         如果参数{@code charset}是{@code null}。
     * @see String#getBytes()
     * @see #sha256(byte[])
     */
    public static byte[] sha256(
            String s,
            Charset charset
    ) {
        rejectIfNull(charset, "charset");

        if (s == null) {
            return sha256((byte[]) null);
        }

        return sha256(s.getBytes(charset));
    }

    /**
     * 使用{@literal MD5}算法对指定的字符串进行哈希。
     * <p>
     * 首先将字符串按照{@literal UTF-8}编码转化为字节数组, 然后进行哈希。
     * </p>
     *
     * @param s
     *         待哈希的字符串。
     * @return 哈希的结果。
     * @see #md5(String, Charset)
     */
    public static byte[] md5(
            String s
    ) {
        return md5(s, UTF_8);
    }

    /**
     * 使用{@literal SHA-1}算法对指定的字符串进行哈希。
     * <p>
     * 首先将字符串按照{@literal UTF-8}编码转化为字节数组, 然后进行哈希。
     * </p>
     *
     * @param s
     *         待哈希的字符串。
     * @return 哈希的结果。
     * @see #sha1(String, Charset)
     */
    public static byte[] sha1(
            String s
    ) {
        return sha1(s, UTF_8);
    }

    /**
     * 使用{@literal SHA-256}算法对指定的字符串进行哈希。
     * <p>
     * 首先将字符串按照{@literal UTF-8}编码转化为字节数组, 然后进行哈希。
     * </p>
     *
     * @param s
     *         待哈希的字符串。
     * @return 哈希的结果。
     * @see #sha256(String, Charset)
     */
    public static byte[] sha256(
            String s
    ) {
        return sha256(s, UTF_8);
    }
}
